package com.itcam.otcbackend.manager.sharding;

import com.baomidou.mybatisplus.extension.toolkit.SqlRunner;
import com.itcam.otcbackend.enums.SpaceLevelEnum;
import com.itcam.otcbackend.enums.SpaceTypeEnum;
import com.itcam.otcbackend.service.SpaceService;
import lombok.extern.slf4j.Slf4j;
import org.apache.shardingsphere.driver.jdbc.core.connection.ShardingSphereConnection;
import org.apache.shardingsphere.infra.metadata.database.rule.ShardingSphereRuleMetaData;
import org.apache.shardingsphere.mode.manager.ContextManager;
import org.apache.shardingsphere.sharding.api.config.ShardingRuleConfiguration;
import org.apache.shardingsphere.sharding.api.config.rule.ShardingTableRuleConfiguration;
import org.apache.shardingsphere.sharding.rule.ShardingRule;
import org.springframework.stereotype.Component;

import javax.annotation.PostConstruct;
import javax.annotation.Resource;
import javax.sql.DataSource;
import java.sql.SQLException;
import java.util.Collections;
import java.util.List;
import java.util.Optional;
import java.util.Set;
import java.util.stream.Collectors;

import com.itcam.otcbackend.model.po.Space;

/**
 * 分表管理器
 */

//@Component
@Slf4j
public class DynamicShardingManager {

    @Resource
    private DataSource dataSource;

    @Resource
    private SpaceService spaceService;

    private static final String LOGIC_TABLE_NAME = "picture";

    private static final String DATABASE_NAME = "logic_db"; // 配置文件中的数据库名称

    @PostConstruct
    public void initialize() {
        log.info("初始化动态分表配置...");
        updateShardingTableNodes();
    }

    /*
        todo
            1. 将管理器注册为 Bean，通过 @PostConstruct 注解，在 Bean 加载后获取所有的分表并更新配置。
            2. 编写获取分表列表的方法，从数据库中查询符合要求的空间列表，再补充上逻辑表，就得到了完整的分表列表。
            3. 更新 ShardingSphere 的 actual-data-nodes 动态表名配置。获取到 ShardingSphere 的 ContextManager，找到配置文件中的那条规则进行更新即可。
     */

    /**
     * 获取所有动态表名，包括初始表 picture 和分表 picture_{spaceId}
     */
    private Set<String> fetchAllPictureTableNames() {
        // 为了测试方便，直接对所有团队空间分表（实际上线改为仅对旗舰版生效）
        Set<Long> spaceIds = spaceService.lambdaQuery()
                .eq(Space::getSpaceType, SpaceTypeEnum.TEAM.getValue())
                .list()
                .stream()
                .map(Space::getId)
                .collect(Collectors.toSet());
        Set<String> tableNames = spaceIds.stream()
                .map(spaceId -> LOGIC_TABLE_NAME + "_" + spaceId)
                .collect(Collectors.toSet());
        tableNames.add(LOGIC_TABLE_NAME); // 添加初始逻辑表
        return tableNames;
    }

    /**
     * 更新 ShardingSphere 的 actual-data-nodes 动态表名配置
     */
    private void updateShardingTableNodes() {
        Set<String> tableNames = fetchAllPictureTableNames();
        String newActualDataNodes = tableNames.stream()
                .map(tableName -> "yu_picture." + tableName) // 确保前缀合法
                .collect(Collectors.joining(","));
        log.info("动态分表 actual-data-nodes 配置: {}", newActualDataNodes);

        // 获取 ShardingSphere 的 ContextManager 实例，该实例用于管理分片规则的上下文信息
        ContextManager contextManager = getContextManager();
        // 从 ContextManager 中获取元数据上下文，进而获取指定数据库的规则元数据
        // 这里的规则元数据包含了分片规则等重要配置信息
        ShardingSphereRuleMetaData ruleMetaData = contextManager.getMetaDataContexts()
                .getMetaData()
                .getDatabases()
                .get(DATABASE_NAME)
                .getRuleMetaData();


        // 从规则元数据中查找单个 ShardingRule 实例，使用 Optional 包装以避免空指针异常
        Optional<ShardingRule> shardingRule = ruleMetaData.findSingleRule(ShardingRule.class);
        // 检查 Optional 中是否存在 ShardingRule 实例
        if (shardingRule.isPresent()) {
            // 获取 ShardingRule 实例的配置信息
            ShardingRuleConfiguration ruleConfig = (ShardingRuleConfiguration) shardingRule.get().getConfiguration();
            // 对规则配置中的表规则进行更新操作
            List<ShardingTableRuleConfiguration> updatedRules = ruleConfig.getTables()
                    .stream()
                    // 遍历每个旧的表规则
                    .map(oldTableRule -> {
                        // 检查当前表规则的逻辑表名是否与预定义的逻辑表名一致
                        if (LOGIC_TABLE_NAME.equals(oldTableRule.getLogicTable())) {
                            // 创建一个新的表规则配置，使用预定义的逻辑表名和新的实际数据节点
                            ShardingTableRuleConfiguration newTableRuleConfig = new ShardingTableRuleConfiguration(LOGIC_TABLE_NAME, newActualDataNodes);
                            // 复制旧表规则的数据库分片策略到新表规则
                            newTableRuleConfig.setDatabaseShardingStrategy(oldTableRule.getDatabaseShardingStrategy());
                            // 复制旧表规则的表分片策略到新表规则
                            newTableRuleConfig.setTableShardingStrategy(oldTableRule.getTableShardingStrategy());
                            // 复制旧表规则的主键生成策略到新表规则
                            newTableRuleConfig.setKeyGenerateStrategy(oldTableRule.getKeyGenerateStrategy());
                            // 复制旧表规则的审计策略到新表规则
                            newTableRuleConfig.setAuditStrategy(oldTableRule.getAuditStrategy());
                            // 返回新的表规则配置
                            return newTableRuleConfig;
                        }
                        // 如果逻辑表名不一致，返回旧的表规则配置
                        return oldTableRule;
                    })
                    // 将流中的元素收集到一个列表中
                    .collect(Collectors.toList());

            ruleConfig.setTables(updatedRules);
            contextManager.alterRuleConfiguration(DATABASE_NAME, Collections.singleton(ruleConfig));
            contextManager.reloadDatabase(DATABASE_NAME);
            log.info("动态分表规则更新成功！");
        } else {
            log.error("未找到 ShardingSphere 的分片规则配置，动态分表更新失败。");
        }
    }

    /**
     * 获取 ShardingSphere ContextManager
     */
    private ContextManager getContextManager() {
        try (ShardingSphereConnection connection = dataSource.getConnection().unwrap(ShardingSphereConnection.class)) {
            return connection.getContextManager();
        } catch (SQLException e) {
            throw new RuntimeException("获取 ShardingSphere ContextManager 失败", e);
        }
    }

    /**
     * 为指定的空间动态创建分表。
     * 仅当空间类型为团队空间且空间级别为旗舰版时，才会创建分表。
     * 创建分表后，会更新 ShardingSphere 的分片规则以包含新表。
     *
     * @param space 要为其创建分表的空间对象
     */
    public void createSpacePictureTable(Space space) {
        // 动态创建分表
        // 仅为旗舰版团队空间创建分表
        if (space.getSpaceType() == SpaceTypeEnum.TEAM.getValue() && space.getSpaceLevel() == SpaceLevelEnum.FLAGSHIP.getValue()) {
            // 获取空间 ID
            Long spaceId = space.getId();
            // 构造新表名
            String tableName = "picture_" + spaceId;
            // 创建新表
            String createTableSql = "CREATE TABLE " + tableName + " LIKE picture";
            try {
                // 执行 SQL 语句创建新表
                SqlRunner.db().update(createTableSql);
                // 更新分表
                updateShardingTableNodes();
            } catch (Exception e) {
                // 记录创建分表失败的日志
                log.error("创建图片空间分表失败，空间 id = {}", space.getId());
            }
        }
    }


}